package com.wutoon.etl.handler;


import com.wutoon.etl.aop.PowerJobLogReport;
import com.wutoon.etl.constant.CommonConstant;
import com.wutoon.etl.pojo.CommonTransfer2DBTask;
import com.wutoon.etl.util.CommonUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicMarkableReference;

/**
 * @author 武兴云/72176468
 * @version 1.0
 * @date 2024/1/22 17:48
 * @description 写入数据
 */
@Slf4j
public class JdbcWriter implements Runnable{
    private final BlockingQueue<Object[]> queue;

    private final CommonTransfer2DBTask task;

    private final Connection conn;

    private final CountDownLatch producerEnd;

    private final CountDownLatch consumerEnd;

    private final AtomicMarkableReference<String> msg;

    private RequestAttributes requestAttributes;

    public JdbcWriter(BlockingQueue<Object[]> queue, CommonTransfer2DBTask task, Connection conn, CountDownLatch producerEnd,
                      CountDownLatch consumerEnd, AtomicMarkableReference<String> msg, RequestAttributes requestAttributes) {
        this.queue = queue;
        this.task = task;
        this.conn = conn;
        this.consumerEnd = consumerEnd;
        this.producerEnd = producerEnd;
        this.msg = msg;
        this.requestAttributes = requestAttributes;
    }

    /**
     * mysql,ob连接串后要加&rewriteBatchedStatements=true，否则会很慢
     * 增量场景目标sql需要时upsert，replace into等
     * 测试：55W数据耗时10S
     */
    @Override
    public void run() {
        log.info("JdbcWriter start......");
        RequestContextHolder.setRequestAttributes(requestAttributes);
        PowerJobLogReport.report("JdbcWriter start......");
        long dataSize = 0;
        long start = System.currentTimeMillis();
        try (PreparedStatement stmt = conn.prepareStatement(task.getTargetSql())) {
            conn.setAutoCommit(false);
            while (true) {
                // 主线程终止了子线程
                if (Thread.currentThread().isInterrupted()) {
                    break;
                }
                Object[] poll = queue.poll(CommonConstant.POLL_TIMEOUT_SECONDS, TimeUnit.SECONDS);
                // 生产者执行结束 并且 队列为空
                if (poll == null && producerEnd.getCount() < CommonConstant.NUMBER_ONE) {
                    break;
                }
                // 获取为空继续下一轮判断
                if (poll == null) {
                    continue;
                }
                // 不同数据库之间etl，需要注意字段类型转换，eg：oracle的timestamp转ob的timestamp，查询时使用cast(xxx as date)
                for (int i = 0; i < poll.length; i++) {
                    stmt.setObject(i + 1, poll[i]);
                }
                stmt.addBatch();
                dataSize ++;
                if (dataSize % task.getBatchInsertSize() == 0) {
                    log.info("JdbcWriter execute batch insert:{}, data size:{}", task.getBatchInsertSize(), dataSize);
                    PowerJobLogReport.report("JdbcWriter execute batch insert:%s, data size:%s", task.getBatchInsertSize(), dataSize);
                    stmt.executeBatch();
                    conn.commit();
                    stmt.clearBatch();
                }
            }
            if (dataSize % task.getBatchInsertSize() != 0) {
                stmt.executeBatch();
                conn.commit();
                stmt.clearBatch();
            }
        } catch (Exception e) {
            log.error("DataTransfer writer error:, ", e);
            // 错误信息上抛
            String stackTraceString = CommonUtil.getStackTraceString(e);
            msg.set(stackTraceString, true);
        } finally {
            consumerEnd.countDown();
            PowerJobLogReport.report("DataTransfer writer end duration:%s, dataSize:%s", System.currentTimeMillis() - start, dataSize);
            log.info("DataTransfer writer end duration:{}, dataSize:{}", System.currentTimeMillis() - start, dataSize);
            RequestContextHolder.resetRequestAttributes();
        }
    }
}
